package com.yc.config;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.serializer.StringRedisSerializer;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;

@Component
public class PersistentMemoryStore implements ChatMemoryStore {
    private static final String REDIS_KEY_PREFIX = "chat:memory:";
    private final RedisTemplate<String, String> redisTemplate;

    @Autowired
    public PersistentMemoryStore(RedisTemplate<String, String> redisTemplate) {
        this.redisTemplate = redisTemplate;
        if (!(redisTemplate.getKeySerializer() instanceof StringRedisSerializer)) {
            throw new IllegalStateException("RedisTemplate must use StringRedisSerializer");
        }
    }

    private String getKey(Object memoryId) {
        return REDIS_KEY_PREFIX + memoryId;
    }

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        String key = getKey(memoryId);
        String messageJson = redisTemplate.opsForValue().get(key);
        if(messageJson == null || messageJson.isEmpty()){
            return List.of();
        }
        return ChatMessageDeserializer.messagesFromJson(messageJson);
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        String key = getKey(memoryId);
        String messageJson = ChatMessageSerializer.messagesToJson(messages);
        redisTemplate.opsForValue().set(key, messageJson);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        String key = getKey(memoryId);
        redisTemplate.delete(key);
    }

    // 获取匹配模式的key（不包含前缀）
    public Set<String> getKeysByPattern(String pattern) {
        Set<String> keysWithPrefix = redisTemplate.keys(REDIS_KEY_PREFIX + pattern);
        Set<String> keysWithoutPrefix = new java.util.HashSet<>();
        for (String key : keysWithPrefix) {
            keysWithoutPrefix.add(key.substring(REDIS_KEY_PREFIX.length()));
        }
        return keysWithoutPrefix;
    }

    // 获取所有匹配模式的消息
    public List<ChatMessage> getAllMessagesByPattern(String pattern) {
        Set<String> keys = getKeysByPattern(pattern);
        List<ChatMessage> allMessages = new ArrayList<>();
        for (String key : keys) {
            allMessages.addAll(getMessages(key));
        }
        return allMessages;
    }

    /**
     * 删除三角色模式的聊天数据（listener/coach/mentor）
     * @param memoryId 基础 memoryId（不含角色后缀）
     */
    public void deleteTripleRoleData(String memoryId) {
        String[] roles = {"listener", "coach", "mentor"};
        for (String role : roles) {
            // 使用getKey方法确保包含前缀
            redisTemplate.delete(getKey(memoryId + "_" + role));
        }
    }

    /**
     * 删除角色互换模式的聊天数据（patient）
     * @param memoryId 基础 memoryId（不含角色后缀）
     */
    public void deletePatientRoleData(String memoryId) {
        // 使用getKey方法确保包含前缀
        redisTemplate.delete(getKey(memoryId + "_patient"));
    }
}